import faiss
import json
import numpy as np
import os

from text2embedding import get_embedding


class ModelEmbeddingSearch:
    def __init__(self, index_file_path="model_index.faiss", model_names_file="model_names.json"):
        self.index_file_path = index_file_path
        self.model_names_file = model_names_file
        if os.path.exists(self.index_file_path):
            self.index = faiss.read_index(self.index_file_path)
        else:
            self.index = faiss.IndexFlatL2(1024)
        if os.path.exists(self.model_names_file):
            with open(self.model_names_file, 'r') as f:
                self.model_names = json.load(f)
        else:
            self.model_names = []

    def insert_models(self, models):
        for model in models:
            prompt = f"{model['type']}, {model['description']}"
            embedding = get_embedding(prompt)
            self.index.add(embedding.reshape(1, -1))
            self.model_names.append(model['name'])
        faiss.write_index(self.index, self.index_file_path)
        with open(self.model_names_file, 'w') as f:
            json.dump(self.model_names, f)

    def search_model(self, text):
        embedding = get_embedding(text)
        _, indices = self.index.search(embedding.reshape(1, -1), 1)
        index = indices[0][0]
        return self.model_names[index]


if __name__ == "__main__":
    # 示例的 model-list.json 数据
    # 从json文件中读取模型列表
    with open("model-list.json", 'r') as f:
        model_list = json.load(f)

    searcher = ModelEmbeddingSearch()
    # searcher.insert_models(model_list["models"])
    
    query_text = "You can use this model in the [a/ComfyUI IPAdapter plus](https://github.com/cubiq/ComfyUI_IPAdapter_plus) extension."
    result = searcher.search_model(query_text)
    print(f"最相似的模型名称: {result}")
    